import matplotlib.pyplot as plt
import matplotlib.ticker as tck
import numpy as np
from scipy.optimize import leastsq
import operator
import os
from os.path import join

PATH = os.getcwd() 

def lorentzian(lor_x,lor):
	numerator =  lor[2]*(lor[0]**2 ) - 2*lor[3]*lor[0]*( lor_x- (lor[1]) )
	# numerator =  lor[2]*(lor[0]**2 )
	denominator = ( lor_x- (lor[1]) )**2 + lor[0]**2
	lor_y = (numerator/denominator) + lor[4]
	return lor_y

def lorentzian_sym(lor_x,lor):
	numerator =  lor[2]*(lor[0]**2 )
	denominator = ( lor_x- (lor[1]) )**2 + lor[0]**2
	lor_y = (numerator/denominator) + lor[4]
	return lor_y

def residuals(p,y,x):
	err = y - lorentzian(x,p)
	return err

rem_points = 25 # number of removed points near zero field

fileList = [f for f in os.listdir(PATH) if f.split('.')[-1] == 'csv']

outFile = open('all_param.txt', 'a')
outFile.write('Freq, Linewidth, Reso_field, A1, A2 \n')

for inFile in fileList:
	datafile = join(PATH, inFile)
	data = np.loadtxt(datafile, delimiter = ',', comments = 'F')
	
	Reso_freq = float((inFile.split('-')[2]).split('G')[0])

	for index in [1,2]:

		
		if index == 2:
			xdata = data[len(data)/2+rem_points:, 1]# a[x,y] means daata of array a at row x and column y
			ydata = 1e6*data[len(data)/2+rem_points:, 2]
			peak_index, peak_value = min(enumerate(ydata), key=operator.itemgetter(1))
		if index == 1:
			xdata = data[:len(data)/2-rem_points, 1]
			ydata = 1e6*data[:len(data)/2-rem_points, 2]
			peak_index, peak_value = max(enumerate(ydata), key=operator.itemgetter(1))

		
		alpha = 0.01
		p0 = [alpha, xdata[peak_index], 0.7*(min(ydata)-max(ydata)), 0.3*(min(ydata)-max(ydata)), max(ydata)]

		
		pbest = leastsq(residuals,p0,args=(ydata,xdata),full_output=1)

		best_parameters = pbest[0]
		print (best_parameters)
		SiGn = best_parameters[0]/abs(best_parameters[0])
		
		outFile.write(str(Reso_freq) + ',' + str(best_parameters[0]) + ',' + str(best_parameters[1]) + ',' + str(best_parameters[2]) + ',' + str(2*best_parameters[3]) + '\n')
		# plot data and fit together

		plt.figure(figsize = (6,3))
		plt.plot(np.array(xdata), ydata, 'o', ms= 4)        
		fit_x = np.linspace(min(xdata), max(xdata), 1000)
		fit_y = lorentzian(fit_x, best_parameters)
		fity1 = lorentzian_sym(fit_x, best_parameters)
		fity2 = lorentzian(fit_x, best_parameters) - lorentzian_sym(fit_x, best_parameters)
		fit_file = open(inFile.split('.')[0] + str(2*index) + '.txt', 'a')
		fit_file.write('fit_x\tfit_y\n')
		for x,y in zip(fit_x, fit_y):
			fit_file.write(str(x) + '\t' + str(y))
			fit_file.write('\n')
		fit_file.close()

		plt.plot(fit_x, fit_y, '-',lw=2,color='red')

		plt.xlabel('Magnetic field (kOe)')
		plt.ylabel(r'ST-FMR signal ($\mu$V)')
		plt.subplots_adjust(bottom = 0.17)
		plt.subplots_adjust(right = 0.75)

		plt.figtext(0.77,0.5,'$A_1$ = '+str(round(best_parameters[2],2))+' '+r'$\mu$V', fontdict=None)
		plt.figtext(0.77,0.4,'$A_2$ = '+str(round(SiGn*best_parameters[3],2))+' '+r'$\mu$V', fontdict=None)
		plt.figtext(0.77,0.3,'$B_0$ = '+str(round(best_parameters[1],2))+' '+'kOe', fontdict=None)
		lg = plt.legend(('$meas$', '$fit$'), prop={'size':15}, labelspacing=0.2, fancybox = True, bbox_to_anchor = (1.4, 1))
		

		ax = plt.gca()
		ax.xaxis.set_major_locator(tck.MaxNLocator(5))
		ax.yaxis.set_major_locator(tck.MaxNLocator(5))


		plt.savefig(inFile.split('.')[0]+'for fitting'+str(2*index-1)+'.png', dpi = 300)

		plt.figure(figsize = (5.5,3.3))
		plt.plot(fit_x, fity1 - best_parameters[4] , '*')
		plt.plot(fit_x, fity2, 's')
		lg = plt.legend(('$Symmetric$', '$Asymmetric$'), prop={'size':12}, labelspacing=0.2, fancybox = True, loc = 0)
		plt.show()
		plt.savefig(inFile.split('.')[0]+'for fitting'+str(2*index)+'.png', dpi = 300)

outFile.close()